package com.wj.tree;

import com.wj.linear.Stack;

import java.io.*;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author wen.jie
 * @date 2021/8/26 16:33
 */
public class HuffmanCode {

    private Node root = null;

    //赫夫曼编码表
    private Map<Byte, String> huffmanCodes = new HashMap<>();
    private StringBuilder sb = new StringBuilder();
    private String content;
    private final String LEFT = "0";
    private final String RIGHT = "1";

    public HuffmanCode(String content){
        this.content = content;
        //创建节点
        List<Node> nodes = createNodes(content.getBytes());
        //创建赫夫曼树
        root = createHuffmanTree(nodes);
        getCodes(root);
    }

    public HuffmanCode(){}

    public void zipFile(String src, String dest) {
        try (FileInputStream fis = new FileInputStream(src);
             ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(dest))) {
            byte[] b = new byte[fis.available()];
            fis.read(b);
            byte[] huffmanBytes = huffmanZip(b);
            oos.writeObject(huffmanBytes);
            oos.writeObject(huffmanCodes);
        }catch (IOException e){
            e.printStackTrace();
        }
    }

    public void unzipFile(String zipFile, String destFile) {
        try (ObjectInputStream ois = new ObjectInputStream(new FileInputStream(zipFile));
             FileOutputStream fos = new FileOutputStream(destFile)){
            byte[] huffmanBytes = (byte[])ois.readObject();
            Map<Byte, String> codes = (Map<Byte, String>)ois.readObject();
            byte[] bytes = decode(codes, huffmanBytes).getBytes();
            fos.write(bytes);
        } catch (IOException | ClassNotFoundException e){
            e.printStackTrace();
        }
    }

    public byte[] huffmanZip(byte[] bytes) {
        List<Node> nodes = createNodes(bytes);
        //创建赫夫曼树
        root = createHuffmanTree(nodes);
        getCodes(root);
        return huffmanZip(bytes, huffmanCodes);
    }

    public byte[] huffmanZip(){
        return huffmanZip(content.getBytes(), huffmanCodes);
    }

    public Map<Byte, String> getHuffmanCodes() {
        return huffmanCodes;
    }

    /**
     * @author wen.jie
     * @date 2021/8/26 20:15
     * @param huffmanCodes 赫夫曼编码表
     * @param huffmanBytes 赫夫曼编码得到的字节数组
     */
    public String decode(Map<Byte, String> huffmanCodes, byte[] huffmanBytes) {
        //先得到huffmanBytes对应的二进制的字符串
        StringBuilder sb = new StringBuilder();
        for (int i = 0; i < huffmanBytes.length; i++) {
            boolean flag = (i == huffmanBytes.length -1);
            sb.append(byteToBitString(huffmanBytes[i], !flag));
        }
        HashMap<String, Byte> map = new HashMap<>();
        for (Byte bt : huffmanCodes.keySet()){
            map.put(huffmanCodes.get(bt), bt);
        }
        ArrayList<Byte> list = new ArrayList<>();
        Stack<String> stack = new Stack<>();
        for (int i = 0; i < sb.length(); i++) {
            String str = sb.substring(i, i + 1);
            stack.push(str);
            String key = jointStack(stack);
            if(map.containsKey(key)){
                list.add(map.get(key));
                stack.removeAll();
            }
        }
        byte[] b = new byte[list.size()];
        for (int i = 0; i < b.length; i++)
            b[i] = list.get(i);
        return new String(b);
    }

    /**
     * @author wen.jie
     * @date 2021/8/26 21:02
     * 拼接栈中所有元素
     */
    private static String jointStack(Stack<String> stack){
        StringBuilder builder = new StringBuilder();
        for (String str : stack) {
            builder.append(str);
        }
        return builder.reverse().toString();
    }

    /**
     * @author wen.jie
     * @date 2021/8/26 19:51
     * 将一个byte转成一个二进制字符串
     * @param flag 标志是否需要补高位,最后一位不需要补高位
     */
    private static String byteToBitString(byte b, boolean flag) {
        return flag
                ? Integer.toBinaryString((b & 0xFF) + 0x100).substring(1)
                : Integer.toBinaryString(b);
    }

    private static byte[] huffmanZip(byte[] bytes, Map<Byte, String> huffmanCodes) {
        StringBuilder stringBuilder = new StringBuilder();
        for (byte b : bytes)
            stringBuilder.append(huffmanCodes.get(b));
        int len = (stringBuilder.length() + 7) / 8;
        byte[] huffmanCodeBytes = new byte[len];
        //每8位对应一个byte,所以步长为8
        int index = 0;
        for (int i = 0; i < stringBuilder.length(); i+= 8) {
            String strByte;
            if(i+8 > stringBuilder.length()){
                strByte = stringBuilder.substring(i);
            }else {
                strByte = stringBuilder.substring(i, i +8);
            }
            //将strByte转成一个byte
            huffmanCodeBytes[index++] = (byte)Integer.parseInt(strByte, 2);
        }
        return huffmanCodeBytes;
    }

    private void getCodes(Node root){
        if (root == null) return;
        //如果只有一个根节点,就直接放进huffmanCodes中
        if (root.left == null && root.right == null)
            huffmanCodes.put(root.data, LEFT);
        getCodes(root.left, LEFT, sb);
        getCodes(root.right, RIGHT, sb);
    }

    /**
     * @author wen.jie
     * @date 2021/8/26 17:13
     * @param code 左子节点为0,右子节点为1
     */
    private void getCodes(Node node, String code, StringBuilder stringBuilder){
        StringBuilder builder = new StringBuilder(stringBuilder);
        builder.append(code);
        if(node != null ) {
            if(node.data == null){
                //非叶子节点
                getCodes(node.left, LEFT, builder);
                getCodes(node.right, RIGHT, builder);
            }else {
                //叶子节点
                huffmanCodes.put(node.data, builder.toString());
            }
        }
    }

    private void preOrder(){
        if(root == null) return;
        //前序遍历
        root.preOrder();
    }

    private Node createHuffmanTree(List<Node> nodes) {
        while (nodes.size() > 1) {
            Collections.sort(nodes);
            Node leftNode = nodes.get(0);
            Node rightNode = nodes.get(1);
            Node parent = new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;
            nodes.remove(leftNode);
            nodes.remove(rightNode);
            nodes.add(parent);
        }
        return nodes.get(0);
    }

    private List<Node> createNodes(byte[] contentBytes ){
        Byte[] bytes = new Byte[contentBytes.length];
        for (int i = 0; i < contentBytes.length; i++) bytes[i] = contentBytes[i];
        ArrayList<Node> nodes = new ArrayList<>();
        Map<Byte, Integer> collect = Arrays.stream(bytes)
                .collect(Collectors.groupingBy(Function.identity(), Collectors.reducing(0, e -> 1, Integer::sum)));
        collect.forEach((k, v) -> nodes.add(new Node(k, v)));
        return nodes;
    }

    private static class Node implements Comparable<Node>{
        //存放数据本身
        Byte data;
        //权值,表示字符出现的次数
        int weight;
        Node left;
        Node right;

        public Node(Byte data, int weight) {
            this.data = data;
            this.weight = weight;
        }

        //前序遍历
        void preOrder() {
            System.out.println(toString());
            if(this.left != null){
                this.left.preOrder();
            }
            if(this.right != null){
                this.right.preOrder();
            }
        }

        @Override
        public int compareTo(Node o) {
            return this.weight - o.weight;
        }

        @Override
        public String toString() {
            return "Node{" +
                    "data=" + data +
                    ", weight=" + weight +
                    '}';
        }
    }
}
